#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 1;

class KernelScatterSub {
 public:
  __aicore__ inline KernelScatterSub() {}

  __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t blockLength,
                              uint32_t tileNum, uint32_t tileLength,
                              uint32_t lasttileLength, uint32_t stride,
                              uint32_t compute_num, uint32_t magic,
                              uint32_t tilingKey) {
    ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
    
    this->stride = stride;
    this->compute_num = compute_num;
    this->blockLength = blockLength;
    this->magic = magic;
    this->key = tilingKey;
    this->tileNum =
        tileNum ASSERT(tileNum != 0 && "tile num can not be zero!");
    this->tileLength = tileLength / BUFFER_NUM;
    this->lasttileLength = lasttileLength;

    xGm.SetGlobalBuffer((__gm__ DTYPE_VAR*)x + this->blockLength * GetBlockIdx(),
                        this->blockLength);
    yGm.SetGlobalBuffer((__gm__ int32_t*)y + this->blockLength * GetBlockIdx(),
                        (this->compute_num + 31) / 32 * 32);
    zGm.SetGlobalBuffer((__gm__ DTYPE_VAR*)z + this->blockLength * GetBlockIdx(),
                        this->blockLength);            
    outGm.SetGlobalBuffer(
        (__gm__ DTYPE_VAR*)x + this->blockLength * GetBlockIdx(),
        this->blockLength);
    init_x((DTYPE_VAR)0);
  }
  template<typename T>
  __aicore__ inline void init_x(T f) {
    pipe.InitBuffer(inQueueIN, BUFFER_NUM, this->tileLength * 2 * sizeof(T));
    pipe.InitBuffer(outQueueOUT, BUFFER_NUM, this->tileLength * sizeof(T));
  }
  template<>
  __aicore__ inline void init_x(int8_t f) {
    pipe.InitBuffer(inQueueIN, BUFFER_NUM, this->tileLength * 2 * sizeof(int8_t));
    pipe.InitBuffer(outQueueOUT, BUFFER_NUM, this->tileLength * sizeof(int8_t));

    pipe.InitBuffer(tmpBuffer3, this->tileLength * sizeof(half));
    pipe.InitBuffer(tmpBuffer4, this->tileLength * sizeof(half));
    pipe.InitBuffer(tmpBuffer5, this->tileLength * sizeof(int16_t));
    pipe.InitBuffer(tmpBuffer6, this->tileLength * sizeof(int16_t));
  }
  
  __aicore__ inline void Process() {
    for (int32_t index = 0; index < this->compute_num; index++) {
      int32_t loopCount = this->tileNum * BUFFER_NUM;
      uint32_t bios = (uint32_t)yGm.GetValue(index);
      for (int32_t i = 0; i < loopCount; i++) {
        CopyIn(i, bios, (uint32_t)index);
        Compute(i, bios, (uint32_t)index, (DTYPE_VAR)0);
        CopyOut(i, bios, (uint32_t)index);
      }
    }
  }
  
private:
  __aicore__ inline void CopyIn(int32_t progress, uint32_t bios, uint32_t update_index) {
    LocalTensor<DTYPE_VAR> inLocal = inQueueIN.AllocTensor<DTYPE_VAR>();
    if (this->key == 1) {
      DataCopy(inLocal[0], xGm[bios * this->stride], this->lasttileLength);
      DataCopy(inLocal[this->tileLength], zGm[update_index * this->stride], this->lasttileLength);
      for (int32_t index = magic; index < this->lasttileLength; index++) {
        inLocal.SetValue(index, xGm.GetValue(index + bios * this->stride));
        inLocal.SetValue(index + this->tileLength, zGm.GetValue(index + update_index * this->stride));
      }
    } else if (this->key == 2) {
      if (progress == this->tileNum - 1 && this->lasttileLength != 0) {
        DataCopy(inLocal[0], xGm[bios * this->stride + progress * this->tileLength], this->lasttileLength);
        DataCopy(inLocal[this->tileLength], zGm[update_index * this->stride + progress * this->tileLength], this->lasttileLength);
        for (int32_t index = magic; index < this->lasttileLength; index++) {
          inLocal.SetValue(index, xGm.GetValue(bios * this->stride + index + progress * this->tileLength));
          inLocal.SetValue(index + this->tileLength, zGm.GetValue(update_index * this->stride + index + progress * this->tileLength));
        }
      } else {
        DataCopy(inLocal[0], xGm[bios * this->stride + progress * this->tileLength], this->tileLength);
        DataCopy(inLocal[this->tileLength], zGm[update_index * this->stride + progress * this->tileLength], this->tileLength);
      }
    }

    inQueueIN.EnQue(inLocal);
  }

  template<typename T>
  __aicore__ inline void Compute(int32_t progress, uint32_t bios, uint32_t update_index, T f) {
    LocalTensor<T> inLocal = inQueueIN.DeQue<T>();
    LocalTensor<T> xLocal = inLocal;
    LocalTensor<T> zLocal = inLocal[this->tileLength];

    LocalTensor<T> outLocal = outQueueOUT.AllocTensor<T>();

    Sub(outLocal, xLocal, zLocal, this->tileLength);
    
    outQueueOUT.EnQue<T>(outLocal);

    inQueueIN.FreeTensor(inLocal);
  }
  template<>
  __aicore__ inline void Compute(int32_t progress, uint32_t bios, uint32_t update_index, int8_t f) {
    LocalTensor<int8_t> inLocal = inQueueIN.DeQue<int8_t>();
    LocalTensor<int8_t> xLocal = inLocal;
    LocalTensor<int8_t> zLocal = inLocal[this->tileLength];
    LocalTensor<half> temp3 = tmpBuffer3.Get<half>();
    LocalTensor<half> temp4 = tmpBuffer4.Get<half>();
    LocalTensor<int16_t> temp5 = tmpBuffer5.Get<int16_t>();
    LocalTensor<int16_t> temp6 = tmpBuffer6.Get<int16_t>();
    Cast(temp3, xLocal, RoundMode::CAST_NONE, this->tileLength);
    Cast(temp4, zLocal, RoundMode::CAST_NONE, this->tileLength);
    Cast(temp5, temp3, RoundMode::CAST_RINT, this->tileLength);
    Cast(temp6, temp4, RoundMode::CAST_RINT, this->tileLength);
    LocalTensor<int8_t> outLocal = outQueueOUT.AllocTensor<int8_t>();
    
    Sub(temp5, temp5, temp6, this->tileLength);
    Cast(temp3, temp5, RoundMode::CAST_NONE, this->tileLength);
    Cast(outLocal, temp3, RoundMode::CAST_NONE, this->tileLength);
    
    outQueueOUT.EnQue<int8_t>(outLocal);

    inQueueIN.FreeTensor(inLocal);
  }

  __aicore__ inline void CopyOut(int32_t progress, uint32_t bios, uint32_t update_index) {
    LocalTensor<DTYPE_VAR> outLocal = outQueueOUT.DeQue<DTYPE_VAR>();

    if (this->key == 1) {
      DataCopy(outGm[bios * this->stride], outLocal, this->lasttileLength);
      for (int32_t index = magic; index < this->lasttileLength; index++) {
        outGm.SetValue(index + bios * this->stride, outLocal.GetValue(index));
      }
    } else if (this->key == 2) {
      if (progress == this->tileNum - 1 && this->lasttileLength != 0) {
        DataCopy(outGm[bios * this->stride + progress * this->tileLength], outLocal, this->lasttileLength);
        for (int32_t index = magic; index < this->lasttileLength; index++) {
          outGm.SetValue(index + progress * this->tileLength + bios * this->stride, outLocal.GetValue(index));
        }
      } else {
        DataCopy(outGm[progress * this->tileLength + bios * this->stride], outLocal, this->tileLength);
      }
    }

    outQueueOUT.FreeTensor(outLocal);
  }
 private:
  TPipe pipe;
  TQue<QuePosition::VECIN, BUFFER_NUM> inQueueIN;
  TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueOUT;
  GlobalTensor<DTYPE_VAR> xGm;
  GlobalTensor<int32_t> yGm;
  GlobalTensor<DTYPE_VAR> zGm;
  GlobalTensor<DTYPE_VAR> outGm;
  TBuf<> tmpBuffer3;
  TBuf<> tmpBuffer4;
  TBuf<> tmpBuffer5;
  TBuf<> tmpBuffer6;

  uint32_t key;
  uint32_t magic;
  uint32_t compute_num;
  uint32_t stride;
  uint32_t blockLength;
  uint32_t tileNum;
  uint32_t tileLength;
  uint32_t lasttileLength;
};


extern "C" __global__ __aicore__ void scatter_sub(GM_ADDR var, GM_ADDR indices, GM_ADDR updates,  GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelScatterSub op;

    uint32_t tilingKey = 1;
    if (TILING_KEY_IS(1)) {
        tilingKey = 1;
    } else if (TILING_KEY_IS(2)) {
        tilingKey = 2;
    } else {
        tilingKey = 1;
    }

    op.Init(var, indices, updates, tiling_data.blockLength,
            tiling_data.tileNum, tiling_data.tileLength,
            tiling_data.lasttileLength, tiling_data.stride, 
            tiling_data.compute_num, tiling_data.magic,
            tilingKey);
    op.Process();
}
